import os
import yaml
import glob
from typing import Union
from copy import deepcopy as dcp


class BindConf:
    """ 绑定模型与配置

    Args:
        funcs (list): 输入模型类型
        default_conf (str, dict): 输入默认的models中类或函数的配置
    """
    def __init__(self, funcs: list, default_conf: Union[str, dict]=None):
        self.mcfgs = {}
        self.libary = {}
        
        self.models = {}
        for model in funcs:
            name = model.__name__
            self.models[name] = model

        if default_conf: 
            self.bind(default_conf)
    
    def __parse_yaml_conf(self, conf_dict):
        mcfg = self.mcfgs
        for mkey in conf_dict.keys(): # 解析模型名称
            for ckey in conf_dict[mkey].keys(): # 解析配置名称
                # 解析模型的类型
                if mkey not in mcfg.keys():
                    mcfg[mkey] = {}
                
                if ckey in mcfg[mkey].keys():
                    # 只加载一次
                    continue

                # 模型:配置:参数
                mcfg[mkey][ckey] = conf_dict[mkey][ckey]

    def bind(self, conf: Union[str, dict]): 
        """ 绑定配置到模型

        Args:
            conf (Union[str, dict]): 输入配置.yaml路径或.yaml的文件夹或dict

        Raises:
            ValueError: _description_
            ValueError: _description_
        """
        if isinstance(conf, str):
            if os.path.exists(conf):
                if conf.endswith('.yaml'):
                    with open(conf, 'r') as f:
                        conf_dict = yaml.load(f, yaml.FullLoader)
                        self.__parse_yaml_conf(conf_dict)
                elif os.path.isdir(conf):
                    for file_path in glob.glob(os.path.join(conf, "*.yaml")):
                        with open(file_path, 'r') as f:
                            conf_dict = yaml.load(f, yaml.FullLoader)
                            self.__parse_yaml_conf(conf_dict)
                else:
                    raise ValueError("Error: make sure your conf is yaml file or yaml root directory")
        elif isinstance(conf, dict):
            self.__parse_yaml_conf(conf)
        else:
            raise ValueError("Error: make sure your conf is dict or yaml file or yaml root directory")
    
    def build(self, func_name, conf_name, default_args=None):
        """ 建立模型

        Args:
            func_name (str): 输入模型或函数名称
            conf_name (str): 输入配置名称
            default_args (dict, optional): 输入默认的参数部分或全部. Defaults to None.

        Returns:
            object: 输出模型
        """
        if conf_name is None:
            kargs = {}
        else:
            kargs = self.mcfgs[func_name][conf_name]

        args = dcp(kargs) 
        if default_args is not None:
            args.update(default_args)

        if func_name not in self.libary:
            self.libary[func_name] = {}
        if conf_name not in self.libary[func_name].keys():
            self.libary[func_name][conf_name] = args

        model = self.models[func_name]
        return model(**args)
    
    def get_conf(self, func_name, conf_name):
        """ 获取配置

        Args:
            func_name (str): 类别的名称
            conf_name (str): 配置的名称

        Returns:
            dict: 返回配置下的所有参数
        """
        if func_name is None:
            return dcp(self.mcfgs)
        elif conf_name is None:
            return dcp(self.mcfgs[func_name])
        else:
            return dcp(self.mcfgs[func_name][conf_name])

    def project(self):
        """构建项目模型（多个）

        Returns:
            dict: 多个模型，"配置:类型": 模型
        """
        output_models = {}
        for mkey in self.mcfgs.keys():
            for mcfg in self.mcfgs[mkey].keys():
                args = self.mcfgs[mkey][mcfg]
                try:
                    model = self.models[mkey]
                    # print(mkey, model)
                    # print(mkey, mcfg, model, args)
                    output_models[mcfg+":"+mkey] = model(**args)
                except:
                    print("Make sure %s import the project" % mkey)

        return output_models
    
    def save(self, proj_file='./model_config/proj.yaml'):
        """将多个模型配置保存成单个文件

        Args:
            proj_file (str, optional): 输出模型配置位置. Defaults to './model_config/proj.yaml'.
        """
        root = os.path.dirname(proj_file)
        if not os.path.exists(root):
            os.makedirs(root)
        
        with open(proj_file, 'w') as f: 
            yaml.dump(self.libary, f)
